'''
@Company: TWL
@Author: xue jian
@Email: xuejian@kanzhun.com
@Date: 2020-04-23 15:13:01
'''
# -*- coding: UTF-8 -*-
from subprocess import *
import random, time, json, sys
from kafka import KafkaProducer
from kafka import KafkaConsumer
from kafka import TopicPartition


class BatchTrain():
    def __init__(self, model_name, batch_size, shuffle_size, watch_num, neg_rate, process_num, run_shell, net_conf, labels, bootstrap_servers, topic, start_time):
        self.model_name = model_name
        self.batch_size = batch_size
        self.process_num = process_num
        self.shuffle_size = shuffle_size
        self.watch_num = watch_num
        self.neg_rate = neg_rate
        # self.ps_url = ps_url
        self.run_shell = run_shell
        self.net_conf = net_conf
        self.labels = labels
        self.bootstrap_servers = bootstrap_servers
        self.topic = topic
        self.start_time = start_time
        self.shell_p = self.run_shell + " -log /data2/bin/flash/conf/log.properties" + " -nets /data2/bin/flash/conf/" + self.net_conf + " -mn " + self.model_name + " -w " + str(self.watch_num) + " -s " + str(self.shuffle_size) + " -b " + str(self.batch_size) + " -ne " + str(self.neg_rate) + " -log_out /data2/bin/galaxy/log -is_json true -labels " + self.labels
        print("shell = ", self.shell_p)


    def train(self):
        model_name = self.model_name + '_' + str(int(time.time()))
        print('model_name = ', model_name)
        consumer = KafkaConsumer(self.topic, auto_offset_reset='latest', group_id=model_name, bootstrap_servers=self.bootstrap_servers)
        # consumer.unsubscribe()

        # partitions = consumer.partitions_for_topic(topic)
        # time_dict = {}
        # for part in partitions:
        #     time_dict[TopicPartition(topic, part)] = self.start_time

        # ot_dict = consumer.offsets_for_times(time_dict)
        # print(ot_dict)
        # for k, v in ot_dict.items():
        #     if v != 'null':
        #         consumer.assign([k])
        #         consumer.seek(k, v[0])
        #         print(consumer.end_offsets([k]))
        #         last_offset = consumer.end_offsets([k])[k]
        #         no_consume_num = last_offset - v[0]
        #         print('no_consume_num = ', no_consume_num)
        
        rt_proc = Popen(self.shell_p, stdin=PIPE, shell=True)
        rt_in = rt_proc.stdin
        i = 0
        count = 0
        for msg in consumer:
            # if (count < 3):
            #     count += 1
            #     print('message = ', msg)
            #     continue
            # else:
            #     break
            rt_in.write((msg.value.decode() + '\n').encode(encoding='utf-8'))
            rt_in.flush()
            if i > 100000:
                tp = TopicPartition(msg.topic, msg.partition)
                highwater = consumer.highwater(tp)
                # print(highwater)
                # print(msg.offset)
                lag = (highwater - 1) - msg.offset
                if lag > 50000:
                    timeArray = time.localtime(time.time())
                    otherStyleTime = time.strftime("%Y-%m-%d %H:%M:%S", timeArray)
                    print(str(otherStyleTime) + '; lag = ' +  str(lag))
                i = 0
            i += 1
            


if __name__ == '__main__':
    bootstrap_servers=['172.21.32.178:9092', '172.21.32.154:9092', '172.21.32.41:9092', '172.21.32.50:9092', '172.21.32.98:9092']
    topic = 'boss.arc.recommender.bossrec_fid_flow'
    dates = []
    # dates.extend(["2019-08-" + str(i) for i in range(12, 32)])
    # dates.extend(["2019-09-0" + str(i) for i in range(1, 7)])
    start_time = time.time()*1000  - 5 * 3600 * 1000
    # ps_ip = "172.18.39.202"
    batch_train = BatchTrain(str(sys.argv[1]), 500, 10000, 100000, 0.72, 12, '/data2/bin/flash/flash_train', 'base_recall.net', 'l0,one', bootstrap_servers, topic, start_time)
    batch_train.train()